import Koa from 'koa'
import Router from 'koa-router'
import json from 'koa-json'
import bodyParser from 'koa-bodyparser'
import cors from '@koa/cors'
import fs from 'fs'
import {fileURLToPath} from "url"
import path from "path"
import koa_static from 'koa-static'
import models from './models.json' assert { type: 'json' }
import llm from './llm.js'

const __dirname = path.dirname(fileURLToPath(import.meta.url))
const app = new Koa()

app.use(cors())
app.use(koa_static(__dirname + '/public', { maxAge: 7 * 24 * 60 * 60 * 1000 }))
app.use(bodyParser({ enableTypes: ['json', 'form', 'text'] }))
app.use(json())

const router = new Router()

// 获取模型列表接口
router.get('/v1/models', async (ctx) => {
  ctx.body = {
    "object": "list",
    "data": models.map(m => {
      return { id: m.filename, name: m.filename, exist: fs.existsSync(`./models/${m.filename}`) }
    })
  }
})

// 模型切换接口
router.post('/v1/change_model', async (ctx) => {
  await llm.load(ctx.request.body.model)
  ctx.body = {}
})

// 渲染 HTML 页面
router.get('/c/:id', async (ctx) => {
  ctx.type = 'html'
  ctx.body = fs.createReadStream(__dirname + "/public/index.html")
})

// 聊天接口
router.post('/v1/chat/completions', async (ctx) => {
  ctx.set('Content-Type', 'text/event-stream')
  ctx.set('Cache-Control', 'no-cache')
  ctx.set('Connection', 'keep-alive')

  ctx.request.body.messages.forEach(element => {
    element.content = element.content[0].text
  })

  ctx.body = await llm.complete(ctx.request.body.id, ctx.request.body.model, ctx.request.body.messages)
})

app.use(router.routes())
app.use(router.allowedMethods())

process.on("uncaughtException", function (err) { console.error(err, "Uncaught exception") })
process.on("unhandledRejection", (reason, promise) => { console.error({ promise, reason }, "Unhandled Rejection at: Promise") })

// 启动服务器
app.listen(1337, () => {
  console.log(`Server is running at http://localhost:1337`)
})
